import { SSEServerTransport } from '@modelcontextprotocol/sdk/server/sse.js'
import express from 'express'
import { createServer } from './chart.js'
import * as dotenv from 'dotenv'

console.error('[Server] Starting SSE server...')

dotenv.config({ quiet: true })

const app = express()

const transports: Map<string, SSEServerTransport> = new Map<
  string,
  SSEServerTransport
>()

app.get('/sse', async (req, res) => {
  let transport: SSEServerTransport
  const { server, cleanup, startNotification } = createServer()

  if (req?.query?.sessionId) {
    const sessionId = req?.query?.sessionId as string
    transport = transports.get(sessionId) as SSEServerTransport
    console.error(
      "[Client] Reconnecting? This shouldn't happen; when client has a sessionId, GET /sse should not be called again.",
      transport.sessionId
    )
  } else {
    transport = new SSEServerTransport('/message', res)
    transports.set(transport.sessionId, transport)

    await server.connect(transport)
    console.error('[Client] Connected: ', transport.sessionId)

    startNotification()

    server.onclose = async () => {
      console.error('[Client] Disconnected: ', transport.sessionId)
      transports.delete(transport.sessionId)
      await cleanup()
    }
  }
})

app.post('/message', async (req, res) => {
  const sessionId = req?.query?.sessionId as string
  const transport = transports.get(sessionId)
  if (transport) {
    console.log('[Client] Message from', sessionId)
    await transport.handlePostMessage(req, res)
  } else {
    console.error(`No transport found for sessionId ${sessionId}`)
  }
})

const PORT = process.env.PORT || 3320

app.listen(PORT, () => {
  console.error(`[Server] Running on port ${PORT}`)
})
